#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO数据集标注工具 - GUI启动函数
提供函数接口，可以传入文件夹路径启动标注工具
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, simpledialog
from PIL import Image, ImageTk
import os
import json
import math
import yaml
from typing import List, Tuple, Dict, Optional


class YOLOAnnotator:
    def __init__(self, root, initial_folder=None):
        self.root = root
        self.root.title("密涅瓦yolo数据标注工具")
        self.root.geometry("1500x900")
        
        # 数据存储
        self.current_image_path = None
        self.current_image = None
        self.photo = None
        self.annotations = []  # 存储当前图像的标注
        self.classes = []  # 空类别列表，用户需要手动添加
        self.negative_class = "负样本"  # 负样本类别
        self.current_class = 0
        self.drawing = False
        self.start_x = 0
        self.start_y = 0
        self.current_bbox = None
        
        # 多段线标注相关
        self.annotation_mode = "rectangle"  # "rectangle" 或 "polygon"
        self.polygon_points = []  # 存储多边形点
        self.current_polygon = None  # 当前正在绘制的多边形
        
        # 撤销/重做功能
        self.undo_stack = []  # 撤销栈
        self.redo_stack = []  # 重做栈
        self.max_undo_steps = 50  # 最大撤销步数
        
        # 标签颜色配置
        self.class_colors = [
            "#FF0000",  # 红色
            "#00FF00",  # 绿色
            "#0000FF",  # 蓝色
            "#FFFF00",  # 黄色
            "#FF00FF",  # 洋红色
            "#00FFFF",  # 青色
            "#FFA500",  # 橙色
            "#800080",  # 紫色
            "#FFC0CB",  # 粉色
            "#A52A2A",  # 棕色
        ]
        
        # 文件夹相关
        self.image_folder = initial_folder
        self.image_files = []
        self.current_image_index = 0
        
        # 创建界面
        self.create_widgets()
        self.setup_bindings()
        
        # 设置窗口关闭事件处理
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        # 如果提供了初始文件夹，自动加载
        if initial_folder and os.path.exists(initial_folder):
            self.create_dataset_structure(initial_folder)
            self.load_folder(initial_folder)
            # 显示欢迎提示
            self.show_welcome_message()
        
    def create_widgets(self):
        """创建GUI组件"""
        # 主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 左侧控制面板 - 添加滚动条
        control_frame = ttk.Frame(main_frame)
        control_frame.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 5))
        
        # 创建滚动条
        canvas = tk.Canvas(control_frame, width=400)
        scrollbar = ttk.Scrollbar(control_frame, orient="vertical", command=canvas.yview)
        scrollable_frame = ttk.Frame(canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        # 将原来的control_frame改为scrollable_frame
        control_frame = scrollable_frame
        
        # 文件操作
        file_frame = ttk.LabelFrame(control_frame, text="文件操作")
        file_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Button(file_frame, text="打开图像", command=self.open_image).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="打开文件夹", command=self.open_folder).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="保存标注", command=self.save_annotations).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="加载标注", command=self.load_annotations).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="验证数据", command=self.validate_dataset).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="完成标注并检查", command=self.complete_and_check).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="标注指南", command=self.show_annotation_guide).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="导出数据", command=self.export_dataset).pack(fill=tk.X, pady=2)
        
        
        # 文件夹导航
        nav_frame = ttk.LabelFrame(control_frame, text="文件夹导航")
        nav_frame.pack(fill=tk.X, pady=(0, 5))
        
        nav_buttons_frame = ttk.Frame(nav_frame)
        nav_buttons_frame.pack(fill=tk.X, pady=2)
        
        ttk.Button(nav_buttons_frame, text="上一张", command=self.prev_image).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(nav_buttons_frame, text="下一张", command=self.next_image).pack(side=tk.RIGHT, padx=(5, 0))
        
        # 图像信息
        self.image_info_var = tk.StringVar(value="未加载图像")
        info_label = ttk.Label(nav_frame, textvariable=self.image_info_var, wraplength=350)
        info_label.pack(pady=2, padx=5, fill=tk.X)
        
        # 类别管理
        class_frame = ttk.LabelFrame(control_frame, text="类别管理")
        class_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 类别选择
        ttk.Label(class_frame, text="当前类别:").pack(anchor=tk.W)
        self.class_var = tk.StringVar(value="请先添加类别")
        self.class_combo = ttk.Combobox(class_frame, textvariable=self.class_var, 
                                       values=self.classes, state="readonly")
        self.class_combo.pack(fill=tk.X, pady=2)
        self.class_combo.bind("<<ComboboxSelected>>", self.on_class_change)
        
        # 添加类别
        add_class_frame = ttk.Frame(class_frame)
        add_class_frame.pack(fill=tk.X, pady=2)
        self.new_class_var = tk.StringVar()
        ttk.Entry(add_class_frame, textvariable=self.new_class_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        ttk.Button(add_class_frame, text="添加", command=self.add_class).pack(side=tk.RIGHT, padx=(2, 0))
        
        # 负样本处理
        negative_frame = ttk.Frame(class_frame)
        negative_frame.pack(fill=tk.X, pady=2)
        ttk.Button(negative_frame, text="标记为负样本", command=self.mark_as_negative).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(negative_frame, text="取消负样本", command=self.unmark_negative).pack(side=tk.RIGHT, padx=(5, 0))
        
        # 标注质量控制
        quality_frame = ttk.LabelFrame(control_frame, text="标注质量控制")
        quality_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 统计信息
        self.stats_var = tk.StringVar(value="总标注: 0")
        ttk.Label(quality_frame, textvariable=self.stats_var).pack(anchor=tk.W)
        
        # 质量控制按钮
        quality_buttons_frame = ttk.Frame(quality_frame)
        quality_buttons_frame.pack(fill=tk.X, pady=2)
        ttk.Button(quality_buttons_frame, text="检查质量", command=self.check_annotation_quality).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(quality_buttons_frame, text="统计信息", command=self.show_statistics).pack(side=tk.RIGHT, padx=(5, 0))
        
        # 撤销/重做按钮
        undo_redo_frame = ttk.Frame(quality_frame)
        undo_redo_frame.pack(fill=tk.X, pady=2)
        ttk.Button(undo_redo_frame, text="撤销 (Ctrl+Z)", command=self.undo_annotation).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(undo_redo_frame, text="重做 (Ctrl+Y)", command=self.redo_annotation).pack(side=tk.RIGHT, padx=(5, 0))
        
        # 标注模式选择
        mode_frame = ttk.LabelFrame(control_frame, text="标注模式")
        mode_frame.pack(fill=tk.X, pady=(0, 5))
        
        self.mode_var = tk.StringVar(value="rectangle")
        ttk.Radiobutton(mode_frame, text="矩形框", variable=self.mode_var, 
                       value="rectangle", command=self.on_mode_change).pack(anchor=tk.W)
        ttk.Radiobutton(mode_frame, text="多边形", variable=self.mode_var, 
                       value="polygon", command=self.on_mode_change).pack(anchor=tk.W)
        
        # 标注操作
        annotate_frame = ttk.LabelFrame(control_frame, text="标注操作")
        annotate_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Button(annotate_frame, text="确认矩形", command=self.confirm_rectangle).pack(fill=tk.X, pady=2)
        ttk.Button(annotate_frame, text="完成多边形", command=self.finish_polygon).pack(fill=tk.X, pady=2)
        ttk.Button(annotate_frame, text="删除选中", command=self.delete_selected).pack(fill=tk.X, pady=2)
        ttk.Button(annotate_frame, text="清空所有", command=self.clear_annotations).pack(fill=tk.X, pady=2)
        
        # 标注列表
        list_frame = ttk.LabelFrame(control_frame, text="标注列表")
        list_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建Treeview显示标注
        columns = ("ID", "类别", "坐标")
        self.annotation_tree = ttk.Treeview(list_frame, columns=columns, show="headings", height=10)
        
        for col in columns:
            self.annotation_tree.heading(col, text=col)
            self.annotation_tree.column(col, width=80)
        
        # 滚动条
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.annotation_tree.yview)
        self.annotation_tree.configure(yscrollcommand=scrollbar.set)
        
        self.annotation_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 绑定选择事件
        self.annotation_tree.bind("<<TreeviewSelect>>", self.on_annotation_select)
        
        # 右侧图像显示区域
        image_frame = ttk.Frame(main_frame)
        image_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 图像画布
        self.canvas = tk.Canvas(image_frame, bg="white", cursor="crosshair")
        self.canvas.pack(fill=tk.BOTH, expand=True)
        
        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        status_bar = ttk.Label(self.root, textvariable=self.status_var, relief=tk.SUNKEN)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
    def setup_bindings(self):
        """设置键盘和鼠标绑定"""
        # 鼠标事件
        self.canvas.bind("<Button-1>", self.on_mouse_press)
        self.canvas.bind("<B1-Motion>", self.on_mouse_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_release)
        self.canvas.bind("<Button-3>", self.on_right_click)  # 右键完成多边形
        
        # 键盘事件
        self.root.bind("<Key>", self.on_key_press)
        self.root.focus_set()
        
        # 快捷键绑定
        self.root.bind("<Control-z>", lambda e: self.undo_annotation())
        self.root.bind("<Control-y>", lambda e: self.redo_annotation())
        self.root.bind("<Control-s>", lambda e: self.save_annotations())
        self.root.bind("<Control-o>", lambda e: self.open_folder())
        self.root.bind("<Delete>", lambda e: self.delete_selected_annotation())
        self.root.bind("<Escape>", lambda e: self.cancel_current_annotation())
        
        # 数字键快速切换类别
        for i in range(1, 10):
            self.root.bind(f"<Key-{i}>", lambda e, num=i: self.quick_select_class(num-1))
        
    def load_folder(self, folder_path):
        """加载文件夹中的图像"""
        if not os.path.exists(folder_path):
            messagebox.showerror("错误", f"文件夹不存在: {folder_path}")
            return False
            
        self.image_folder = folder_path
        
        # 加载标签信息
        self.load_classes_file()
        
        # 获取所有图像文件
        image_files = [f for f in os.listdir(folder_path) 
                      if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))]
        
        if not image_files:
            messagebox.showwarning("警告", "文件夹中没有找到图像文件")
            return False
        
        # 重命名图像文件为数字编号
        self.rename_images_to_numbers(folder_path, image_files)
        
        # 重新获取重命名后的图像文件列表
        self.image_files = [f for f in os.listdir(folder_path) 
                          if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))]
        
        # 按文件名排序（数字编号）
        self.image_files.sort()
        
        if self.image_files:
            self.current_image_index = 0
            self.load_image(os.path.join(folder_path, self.image_files[0]))
            self.update_image_info()
            self.status_var.set(f"已加载文件夹，共{len(self.image_files)}张图像")
            return True
        else:
            messagebox.showwarning("警告", "重命名后没有找到图像文件")
            return False
            
    def update_image_info(self):
        """更新图像信息显示"""
        if self.image_files:
            current_file = self.image_files[self.current_image_index]
            self.image_info_var.set(f"{self.current_image_index + 1}/{len(self.image_files)}: {current_file}")
        else:
            self.image_info_var.set("未加载图像")
            
    def prev_image(self):
        """上一张图像"""
        if not self.image_files:
            return
            
        # 保存当前标注
        self.save_current_annotations()
        
        if self.current_image_index > 0:
            self.current_image_index -= 1
            self.load_image(os.path.join(self.image_folder, self.image_files[self.current_image_index]))
            self.update_image_info()
            
    def next_image(self):
        """下一张图像"""
        if not self.image_files:
            return
            
        # 保存当前标注
        self.save_current_annotations()
        
        if self.current_image_index < len(self.image_files) - 1:
            self.current_image_index += 1
            self.load_image(os.path.join(self.image_folder, self.image_files[self.current_image_index]))
            self.update_image_info()
            
    def save_current_annotations(self):
        """保存当前图像的标注"""
        if not self.current_image_path or not self.annotations:
            return
            
        try:
            base_name = os.path.splitext(self.current_image_path)[0]
            annotation_file = base_name + ".txt"
            img_width, img_height = self.current_image.size
            
            with open(annotation_file, 'w', encoding='utf-8') as f:
                for annotation in self.annotations:
                    class_id = annotation['class_id']
                    annotation_type = annotation.get('type', 'rectangle')
                    
                    if annotation_type == 'rectangle':
                        bbox = annotation['bbox']
                        x_center = (bbox[0] + bbox[2]) / 2 / img_width
                        y_center = (bbox[1] + bbox[3]) / 2 / img_height
                        width = (bbox[2] - bbox[0]) / img_width
                        height = (bbox[3] - bbox[1]) / img_height
                        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
                        
                    elif annotation_type == 'polygon':
                        points = annotation['points']
                        line = f"{class_id}"
                        for x, y in points:
                            norm_x = x / img_width
                            norm_y = y / img_height
                            line += f" {norm_x:.6f} {norm_y:.6f}"
                        f.write(line + "\n")
        except Exception as e:
            print(f"保存标注时出错: {e}")
    
    def create_dataset_structure(self, folder_path):
        """创建数据集目录结构"""
        try:
            # 创建数据集文件夹
            dataset_folder = os.path.join(folder_path, "数据集")
            os.makedirs(dataset_folder, exist_ok=True)
            
            # 创建子文件夹
            subfolders = ["train", "valid", "test"]
            for subfolder in subfolders:
                subfolder_path = os.path.join(dataset_folder, subfolder)
                os.makedirs(subfolder_path, exist_ok=True)
                # 创建images和labels子文件夹
                os.makedirs(os.path.join(subfolder_path, "images"), exist_ok=True)
                os.makedirs(os.path.join(subfolder_path, "labels"), exist_ok=True)
            
            # 创建data.yaml文件
            yaml_path = os.path.join(dataset_folder, "data.yaml")
            self.create_data_yaml(yaml_path)
            
            # 创建空的标签文件
            classes_file = os.path.join(folder_path, "classes.txt")
            if not os.path.exists(classes_file):
                with open(classes_file, 'w', encoding='utf-8') as f:
                    f.write("")
            
            print(f"已创建数据集目录结构: {dataset_folder}")
            
        except Exception as e:
            print(f"创建数据集目录结构时出错: {e}")
    
    def create_data_yaml(self, yaml_path):
        """创建data.yaml文件"""
        try:
            data = {
                'train': '../train/images',
                'val': '../valid/images', 
                'test': '../test/images',
                'nc': len(self.classes),
                'names': self.classes.copy()
            }
            
            with open(yaml_path, 'w', encoding='utf-8') as f:
                yaml.dump(data, f, default_flow_style=False, allow_unicode=True)
                
        except Exception as e:
            print(f"创建data.yaml文件时出错: {e}")
    
    def save_classes_file(self):
        """保存标签信息到classes.txt文件"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            return
            
        try:
            classes_file = os.path.join(self.image_folder, "classes.txt")
            with open(classes_file, 'w', encoding='utf-8') as f:
                for class_name in self.classes:
                    f.write(f"{class_name}\n")
            print(f"标签信息已保存到: {classes_file}")
        except Exception as e:
            print(f"保存标签文件时出错: {e}")
    
    def load_classes_file(self):
        """从classes.txt文件加载标签信息"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            return
            
        try:
            classes_file = os.path.join(self.image_folder, "classes.txt")
            if os.path.exists(classes_file):
                with open(classes_file, 'r', encoding='utf-8') as f:
                    self.classes = [line.strip() for line in f if line.strip()]
                
                # 更新界面
                self.class_combo['values'] = self.classes
                if self.classes:
                    self.class_var.set(self.classes[0])
                    self.current_class = 0
                else:
                    self.class_var.set("请先添加类别")
                    self.current_class = -1
                
                print(f"已加载标签信息: {len(self.classes)}个标签")
            else:
                print("未找到classes.txt文件，使用空标签列表")
        except Exception as e:
            print(f"加载标签文件时出错: {e}")
            self.classes = []
    
    def update_data_yaml(self):
        """更新data.yaml文件"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            return
            
        try:
            dataset_folder = os.path.join(self.image_folder, "数据集")
            yaml_path = os.path.join(dataset_folder, "data.yaml")
            
            if os.path.exists(yaml_path):
                data = {
                    'train': '../train/images',
                    'val': '../valid/images',
                    'test': '../test/images', 
                    'nc': len(self.classes),
                    'names': self.classes.copy()
                }
                
                with open(yaml_path, 'w', encoding='utf-8') as f:
                    yaml.dump(data, f, default_flow_style=False, allow_unicode=True)
                    
        except Exception as e:
            print(f"更新data.yaml文件时出错: {e}")
    
    # 负样本处理功能
    def mark_as_negative(self):
        """标记当前图像为负样本"""
        if not self.current_image_path:
            messagebox.showwarning("警告", "请先加载图像")
            return
        
        # 创建负样本标注文件
        annotation_file = self.current_image_path.replace('.jpg', '.txt').replace('.jpeg', '.txt').replace('.png', '.txt')
        with open(annotation_file, 'w', encoding='utf-8') as f:
            f.write("")  # 空文件表示负样本
        
        self.status_var.set("已标记为负样本")
    
    def unmark_negative(self):
        """取消负样本标记"""
        if not self.current_image_path:
            messagebox.showwarning("警告", "请先加载图像")
            return
        
        # 删除负样本标注文件
        annotation_file = self.current_image_path.replace('.jpg', '.txt').replace('.jpeg', '.txt').replace('.png', '.txt')
        if os.path.exists(annotation_file):
            os.remove(annotation_file)
        
        self.status_var.set("已取消负样本标记")
    
    
    # 标注质量控制功能
    def check_annotation_quality(self):
        """检查标注质量"""
        if not self.annotations:
            messagebox.showinfo("信息", "当前图像没有标注")
            return
        
        issues = []
        
        # 检查标注框大小
        for i, annotation in enumerate(self.annotations):
            if annotation.get('type') == 'rectangle':
                bbox = annotation['bbox']
                width = bbox[2] - bbox[0]
                height = bbox[3] - bbox[1]
                
                if width < 10 or height < 10:
                    issues.append(f"标注{i+1}: 标注框过小 ({width:.1f}x{height:.1f})")
                elif width > 500 or height > 500:
                    issues.append(f"标注{i+1}: 标注框过大 ({width:.1f}x{height:.1f})")
        
        if issues:
            messagebox.showwarning("质量检查", "发现以下问题:\n" + "\n".join(issues))
        else:
            messagebox.showinfo("质量检查", "标注质量良好！")
    
    def show_statistics(self):
        """显示统计信息"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            messagebox.showwarning("警告", "请先选择图像文件夹")
            return
        
        # 统计所有图像的标注
        total_images = 0
        total_annotations = 0
        class_counts = {}
        
        for img_file in os.listdir(self.image_folder):
            if img_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
                total_images += 1
                
                # 读取标注文件
                label_file = img_file.replace('.jpg', '.txt').replace('.jpeg', '.txt').replace('.png', '.txt')
                label_path = os.path.join(self.image_folder, label_file)
                
                if os.path.exists(label_path):
                    with open(label_path, 'r', encoding='utf-8') as f:
                        lines = f.readlines()
                        for line in lines:
                            line = line.strip()
                            if line:
                                parts = line.split()
                                if parts:
                                    class_id = int(parts[0])
                                    class_name = self.get_class_name(class_id)
                                    class_counts[class_name] = class_counts.get(class_name, 0) + 1
                                    total_annotations += 1
        
        # 显示统计信息
        stats_text = f"总图像数: {total_images}\n总标注数: {total_annotations}\n\n各类别统计:\n"
        for class_name, count in class_counts.items():
            stats_text += f"{class_name}: {count}\n"
        
        messagebox.showinfo("统计信息", stats_text)
    
    
    # 撤销/重做功能
    def save_state(self):
        """保存当前状态到撤销栈"""
        state = {
            'annotations': self.annotations.copy(),
            'image_path': self.current_image_path
        }
        
        self.undo_stack.append(state)
        
        # 限制撤销栈大小
        if len(self.undo_stack) > self.max_undo_steps:
            self.undo_stack.pop(0)
        
        # 清空重做栈
        self.redo_stack.clear()
    
    def undo_annotation(self):
        """撤销操作"""
        if not self.undo_stack:
            messagebox.showinfo("信息", "没有可撤销的操作")
            return
        
        # 保存当前状态到重做栈
        current_state = {
            'annotations': self.annotations.copy(),
            'image_path': self.current_image_path
        }
        self.redo_stack.append(current_state)
        
        # 恢复上一个状态
        state = self.undo_stack.pop()
        self.annotations = state['annotations']
        
        # 更新显示
        self.update_annotation_list()
        self.draw_annotations()
        self.update_stats()
        self.status_var.set("已撤销操作")
    
    def redo_annotation(self):
        """重做操作"""
        if not self.redo_stack:
            messagebox.showinfo("信息", "没有可重做的操作")
            return
        
        # 恢复重做栈中的状态
        state = self.redo_stack.pop()
        self.annotations = state['annotations']
        
        # 保存到撤销栈
        self.undo_stack.append({
            'annotations': self.annotations.copy(),
            'image_path': self.current_image_path
        })
        
        # 更新显示
        self.update_annotation_list()
        self.draw_annotations()
        self.update_stats()
        self.status_var.set("已重做操作")
    
    # 数据验证和导出功能
    def validate_dataset(self, return_details=False):
        """
        验证数据集
        
        Args:
            return_details: 如果为True，返回详细的问题信息字典
            
        Returns:
            如果return_details为True，返回(是否有问题, 问题信息字典)
            问题信息字典格式: {
                'has_issues': bool,
                'issues': list,  # 问题描述列表
                'problem_files': dict,  # {图片文件名: [问题列表]}
                'total_images': int,
                'total_labels': int
            }
        """
        if not hasattr(self, 'image_folder') or not self.image_folder:
            if return_details:
                return True, {'has_issues': True, 'issues': ['请先选择图像文件夹'], 'problem_files': {}, 'total_images': 0, 'total_labels': 0}
            messagebox.showwarning("警告", "请先选择图像文件夹")
            return
        
        try:
            issues = []
            problem_files = {}  # {图片文件名: [问题列表]}
            
            # 获取所有图像文件
            image_files = []
            for img_file in os.listdir(self.image_folder):
                if img_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
                    image_files.append(img_file)
            
            # 检查图像和标注文件匹配
            for img_file in image_files:
                base_name = os.path.splitext(img_file)[0]
                label_file = base_name + ".txt"
                label_path = os.path.join(self.image_folder, label_file)
                
                if not os.path.exists(label_path):
                    issue_msg = f"缺少标注文件: {label_file}"
                    issues.append(issue_msg)
                    if img_file not in problem_files:
                        problem_files[img_file] = []
                    problem_files[img_file].append(issue_msg)
            
            # 检查标注文件格式
            for label_file in os.listdir(self.image_folder):
                if label_file.endswith('.txt') and label_file != 'classes.txt':
                    label_path = os.path.join(self.image_folder, label_file)
                    base_name = os.path.splitext(label_file)[0]
                    
                    # 查找对应的图像文件
                    img_file = None
                    for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
                        potential_img = base_name + ext
                        if potential_img in image_files or os.path.exists(os.path.join(self.image_folder, potential_img)):
                            img_file = potential_img
                            break
                    
                    try:
                        with open(label_path, 'r', encoding='utf-8') as f:
                            for line_num, line in enumerate(f, 1):
                                line = line.strip()
                                if line:
                                    parts = line.split()
                                    
                                    # 检查格式：至少5个部分（类别ID + 至少4个坐标）
                                    if len(parts) < 5:
                                        issue_msg = f"{label_file}:{line_num} 格式错误（至少需要5个部分）"
                                        issues.append(issue_msg)
                                        if img_file and img_file not in problem_files:
                                            problem_files[img_file] = []
                                        if img_file:
                                            problem_files[img_file].append(issue_msg)
                                    elif len(parts) % 2 == 0:
                                        issue_msg = f"{label_file}:{line_num} 格式错误（坐标数量必须是偶数）"
                                        issues.append(issue_msg)
                                        if img_file and img_file not in problem_files:
                                            problem_files[img_file] = []
                                        if img_file:
                                            problem_files[img_file].append(issue_msg)
                                    else:
                                        try:
                                            class_id = int(parts[0])
                                            if class_id < 0:
                                                issue_msg = f"{label_file}:{line_num} 类别ID不能为负数"
                                                issues.append(issue_msg)
                                                if img_file and img_file not in problem_files:
                                                    problem_files[img_file] = []
                                                if img_file:
                                                    problem_files[img_file].append(issue_msg)
                                            elif class_id >= len(self.classes):
                                                issue_msg = f"{label_file}:{line_num} 类别ID超出范围（最大为{len(self.classes)-1}）"
                                                issues.append(issue_msg)
                                                if img_file and img_file not in problem_files:
                                                    problem_files[img_file] = []
                                                if img_file:
                                                    problem_files[img_file].append(issue_msg)
                                            else:
                                                # 检查坐标值是否在0-1范围内
                                                coords = parts[1:]
                                                for i in range(0, len(coords), 2):
                                                    try:
                                                        x = float(coords[i])
                                                        y = float(coords[i+1])
                                                        if x < 0 or x > 1 or y < 0 or y > 1:
                                                            issue_msg = f"{label_file}:{line_num} 坐标超出范围（应在0-1之间）"
                                                            issues.append(issue_msg)
                                                            if img_file and img_file not in problem_files:
                                                                problem_files[img_file] = []
                                                            if img_file:
                                                                problem_files[img_file].append(issue_msg)
                                                            break
                                                    except ValueError:
                                                        issue_msg = f"{label_file}:{line_num} 坐标不是有效数字"
                                                        issues.append(issue_msg)
                                                        if img_file and img_file not in problem_files:
                                                            problem_files[img_file] = []
                                                        if img_file:
                                                            problem_files[img_file].append(issue_msg)
                                                        break
                                        except ValueError:
                                            issue_msg = f"{label_file}:{line_num} 类别ID不是数字"
                                            issues.append(issue_msg)
                                            if img_file and img_file not in problem_files:
                                                problem_files[img_file] = []
                                            if img_file:
                                                problem_files[img_file].append(issue_msg)
                    except Exception as e:
                        issue_msg = f"读取标注文件失败: {label_file} ({str(e)})"
                        issues.append(issue_msg)
                        if img_file and img_file not in problem_files:
                            problem_files[img_file] = []
                        if img_file:
                            problem_files[img_file].append(issue_msg)
            
            total_images = len(image_files)
            total_labels = len([f for f in os.listdir(self.image_folder) 
                              if f.endswith('.txt') and f != 'classes.txt'])
            
            has_issues = len(issues) > 0
            
            if return_details:
                return has_issues, {
                    'has_issues': has_issues,
                    'issues': issues,
                    'problem_files': problem_files,
                    'total_images': total_images,
                    'total_labels': total_labels
                }
            
            if issues:
                # 统计问题类型
                total_issues = len(issues)
                missing_files = len([i for i in issues if "缺少标注文件" in i])
                format_errors = len([i for i in issues if "格式错误" in i])
                class_errors = len([i for i in issues if "类别ID" in i])
                coord_errors = len([i for i in issues if "坐标" in i])
                
                summary = f"验证完成！发现 {total_issues} 个问题：\n"
                if missing_files > 0:
                    summary += f"- 缺少标注文件: {missing_files} 个\n"
                if format_errors > 0:
                    summary += f"- 格式错误: {format_errors} 个\n"
                if class_errors > 0:
                    summary += f"- 类别ID问题: {class_errors} 个\n"
                if coord_errors > 0:
                    summary += f"- 坐标问题: {coord_errors} 个\n"
                
                details = "\n详细问题列表：\n" + "\n".join(issues[:10])
                if len(issues) > 10:
                    details += f"\n... 还有 {len(issues) - 10} 个问题"
                
                messagebox.showwarning("验证结果", summary + details)
            else:
                # 统计成功信息
                success_msg = f"✅ 数据集验证通过！\n\n"
                success_msg += f"📊 统计信息：\n"
                success_msg += f"- 图像文件: {total_images} 个\n"
                success_msg += f"- 标注文件: {total_labels} 个\n"
                success_msg += f"- 类别数量: {len(self.classes)} 个\n"
                success_msg += f"- 类别列表: {', '.join(self.classes)}\n\n"
                success_msg += f"🎉 所有标注文件格式正确，可以用于YOLO训练！"
                
                messagebox.showinfo("验证结果", success_msg)
                
        except Exception as e:
            error_msg = f"验证数据集失败: {str(e)}"
            if return_details:
                return True, {'has_issues': True, 'issues': [error_msg], 'problem_files': {}, 'total_images': 0, 'total_labels': 0}
            messagebox.showerror("错误", error_msg)
    
    def complete_and_check(self):
        """完成标注并检查，如果有问题则打开问题图片，如果没问题则弹出数据集划分对话框"""
        # 首先保存当前标注
        self.save_current_annotations()
        
        if not hasattr(self, 'image_folder') or not self.image_folder:
            messagebox.showwarning("警告", "请先选择图像文件夹")
            return
        
        # 验证数据集
        has_issues, details = self.validate_dataset(return_details=True)
        
        if has_issues:
            # 有问题，显示问题并打开第一张问题图片
            problem_files = details.get('problem_files', {})
            
            if problem_files:
                # 获取第一张有问题的图片
                first_problem_file = list(problem_files.keys())[0]
                problem_img_path = os.path.join(self.image_folder, first_problem_file)
                
                # 构建问题描述
                problems = problem_files[first_problem_file]
                problem_msg = f"发现标注问题！\n\n"
                problem_msg += f"问题图片: {first_problem_file}\n\n"
                problem_msg += "问题列表:\n"
                for i, prob in enumerate(problems[:5], 1):
                    problem_msg += f"{i}. {prob}\n"
                if len(problems) > 5:
                    problem_msg += f"... 还有 {len(problems) - 5} 个问题\n"
                
                problem_msg += f"\n共有 {len(problem_files)} 张图片存在问题。\n"
                problem_msg += "正在打开第一张问题图片，请修复后重新检查。"
                
                messagebox.showwarning("标注检查失败", problem_msg)
                
                # 打开问题图片
                if os.path.exists(problem_img_path):
                    # 找到图片在列表中的索引
                    try:
                        img_index = self.image_files.index(first_problem_file)
                        self.current_image_index = img_index
                        self.load_image(problem_img_path)
                        self.update_image_info()
                        self.status_var.set(f"已切换到问题图片: {first_problem_file}")
                    except ValueError:
                        # 如果不在列表中，直接加载
                        self.load_image(problem_img_path)
                        self.status_var.set(f"已打开问题图片: {first_problem_file}")
                else:
                    messagebox.showerror("错误", f"找不到问题图片: {problem_img_path}")
            else:
                # 有一般性问题但没有具体文件
                issue_msg = "发现以下问题:\n" + "\n".join(details.get('issues', [])[:10])
                if len(details.get('issues', [])) > 10:
                    issue_msg += f"\n... 还有 {len(details['issues']) - 10} 个问题"
                messagebox.showwarning("标注检查失败", issue_msg)
        else:
            # 没有问题，弹出数据集划分对话框
            self.show_dataset_split_dialog(details)
    
    def show_dataset_split_dialog(self, validation_details):
        """显示数据集划分对话框"""
        dialog = tk.Toplevel(self.root)
        dialog.title("数据集划分")
        dialog.geometry("500x400")
        dialog.transient(self.root)
        dialog.grab_set()
        
        total_images = validation_details.get('total_images', 0)
        if total_images == 0:
            messagebox.showwarning("警告", "没有图像可以划分")
            dialog.destroy()
            return
        
        # 主框架
        main_frame = ttk.Frame(dialog, padding=20)
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 标题
        title_label = ttk.Label(main_frame, text="划分数据集", font=("Arial", 14, "bold"))
        title_label.pack(pady=(0, 10))
        
        # 统计信息
        info_label = ttk.Label(main_frame, text=f"总图像数: {total_images} 张", font=("Arial", 10))
        info_label.pack(pady=(0, 20))
        
        # 比例输入框架
        ratio_frame = ttk.LabelFrame(main_frame, text="数据集划分比例 (%)", padding=10)
        ratio_frame.pack(fill=tk.X, pady=(0, 10))
        
        # 训练集
        train_frame = ttk.Frame(ratio_frame)
        train_frame.pack(fill=tk.X, pady=5)
        ttk.Label(train_frame, text="训练集:", width=10).pack(side=tk.LEFT)
        train_var = tk.StringVar(value="70")
        train_entry = ttk.Entry(train_frame, textvariable=train_var, width=15)
        train_entry.pack(side=tk.LEFT, padx=5)
        ttk.Label(train_frame, text="%").pack(side=tk.LEFT)
        
        # 验证集
        valid_frame = ttk.Frame(ratio_frame)
        valid_frame.pack(fill=tk.X, pady=5)
        ttk.Label(valid_frame, text="验证集:", width=10).pack(side=tk.LEFT)
        valid_var = tk.StringVar(value="20")
        valid_entry = ttk.Entry(valid_frame, textvariable=valid_var, width=15)
        valid_entry.pack(side=tk.LEFT, padx=5)
        ttk.Label(valid_frame, text="%").pack(side=tk.LEFT)
        
        # 测试集
        test_frame = ttk.Frame(ratio_frame)
        test_frame.pack(fill=tk.X, pady=5)
        ttk.Label(test_frame, text="测试集:", width=10).pack(side=tk.LEFT)
        test_var = tk.StringVar(value="10")
        test_entry = ttk.Entry(test_frame, textvariable=test_var, width=15)
        test_entry.pack(side=tk.LEFT, padx=5)
        ttk.Label(test_frame, text="%").pack(side=tk.LEFT)
        
        # 预设比例按钮
        preset_frame = ttk.Frame(main_frame)
        preset_frame.pack(fill=tk.X, pady=(0, 10))
        ttk.Label(preset_frame, text="常用比例:").pack(side=tk.LEFT, padx=(0, 5))
        
        def set_ratio(train, valid, test):
            train_var.set(str(train))
            valid_var.set(str(valid))
            test_var.set(str(test))
        
        ttk.Button(preset_frame, text="7:2:1", command=lambda: set_ratio(70, 20, 10)).pack(side=tk.LEFT, padx=2)
        ttk.Button(preset_frame, text="8:1:1", command=lambda: set_ratio(80, 10, 10)).pack(side=tk.LEFT, padx=2)
        ttk.Button(preset_frame, text="6:2:2", command=lambda: set_ratio(60, 20, 20)).pack(side=tk.LEFT, padx=2)
        
        # 预览信息
        preview_var = tk.StringVar()
        preview_label = ttk.Label(main_frame, textvariable=preview_var, font=("Arial", 9), foreground="gray")
        preview_label.pack(pady=(0, 10))
        
        def update_preview():
            try:
                train_pct = float(train_var.get())
                valid_pct = float(valid_var.get())
                test_pct = float(test_var.get())
                total_pct = train_pct + valid_pct + test_pct
                
                if abs(total_pct - 100.0) > 0.01:
                    preview_var.set(f"⚠️ 比例总和: {total_pct:.1f}% (应为100%)")
                    preview_label.config(foreground="red")
                else:
                    train_count = int(total_images * train_pct / 100)
                    valid_count = int(total_images * valid_pct / 100)
                    test_count = total_images - train_count - valid_count
                    
                    preview_var.set(f"预览: 训练集={train_count}张, 验证集={valid_count}张, 测试集={test_count}张")
                    preview_label.config(foreground="green")
            except ValueError:
                preview_var.set("请输入有效的数字")
                preview_label.config(foreground="red")
        
        train_var.trace('w', lambda *args: update_preview())
        valid_var.trace('w', lambda *args: update_preview())
        test_var.trace('w', lambda *args: update_preview())
        update_preview()
        
        # 按钮框架
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(10, 0))
        
        def split_dataset():
            try:
                train_pct = float(train_var.get())
                valid_pct = float(valid_var.get())
                test_pct = float(test_var.get())
                total_pct = train_pct + valid_pct + test_pct
                
                if abs(total_pct - 100.0) > 0.01:
                    messagebox.showerror("错误", f"比例总和必须等于100%，当前为{total_pct:.1f}%")
                    return
                
                if train_pct < 0 or valid_pct < 0 or test_pct < 0:
                    messagebox.showerror("错误", "比例不能为负数")
                    return
                
                dialog.destroy()
                self.split_and_copy_dataset(train_pct, valid_pct, test_pct)
                
            except ValueError:
                messagebox.showerror("错误", "请输入有效的数字")
        
        ttk.Button(button_frame, text="取消", command=dialog.destroy).pack(side=tk.RIGHT, padx=(5, 0))
        ttk.Button(button_frame, text="确定", command=split_dataset).pack(side=tk.RIGHT)
        
        # 居中显示
        dialog.update_idletasks()
        x = (dialog.winfo_screenwidth() // 2) - (dialog.winfo_width() // 2)
        y = (dialog.winfo_screenheight() // 2) - (dialog.winfo_height() // 2)
        dialog.geometry(f"+{x}+{y}")
    
    def split_and_copy_dataset(self, train_pct, valid_pct, test_pct):
        """按照比例划分数据集并复制文件"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            messagebox.showerror("错误", "未选择图像文件夹")
            return
        
        try:
            import shutil
            import random
            
            # 获取所有图像文件
            image_files = []
            for img_file in os.listdir(self.image_folder):
                if img_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
                    image_files.append(img_file)
            
            if not image_files:
                messagebox.showwarning("警告", "文件夹中没有图像文件")
                return
            
            # 创建数据集文件夹
            dataset_folder = os.path.join(self.image_folder, "数据集")
            os.makedirs(dataset_folder, exist_ok=True)
            
            # 创建目录结构
            subfolders = {
                'train': train_pct / 100.0,
                'valid': valid_pct / 100.0,
                'test': test_pct / 100.0
            }
            
            for subfolder in subfolders.keys():
                images_dir = os.path.join(dataset_folder, subfolder, "images")
                labels_dir = os.path.join(dataset_folder, subfolder, "labels")
                os.makedirs(images_dir, exist_ok=True)
                os.makedirs(labels_dir, exist_ok=True)
            
            # 打乱图像顺序
            random.seed(42)  # 固定随机种子，保证可重复
            shuffled_files = image_files.copy()
            random.shuffle(shuffled_files)
            
            # 计算每个集合的数量
            total = len(shuffled_files)
            train_count = int(total * train_pct / 100)
            valid_count = int(total * valid_pct / 100)
            test_count = total - train_count - valid_count
            
            # 分配文件
            train_files = shuffled_files[:train_count]
            valid_files = shuffled_files[train_count:train_count + valid_count]
            test_files = shuffled_files[train_count + valid_count:]
            
            # 复制文件
            def copy_files(file_list, subset_name):
                copied_count = 0
                for img_file in file_list:
                    # 复制图像文件
                    src_img = os.path.join(self.image_folder, img_file)
                    dst_img = os.path.join(dataset_folder, subset_name, "images", img_file)
                    shutil.copy2(src_img, dst_img)
                    
                    # 复制标注文件
                    base_name = os.path.splitext(img_file)[0]
                    label_file = base_name + ".txt"
                    src_label = os.path.join(self.image_folder, label_file)
                    dst_label = os.path.join(dataset_folder, subset_name, "labels", label_file)
                    
                    if os.path.exists(src_label):
                        shutil.copy2(src_label, dst_label)
                        copied_count += 1
                    else:
                        # 如果没有标注文件，创建一个空文件（负样本）
                        with open(dst_label, 'w', encoding='utf-8') as f:
                            f.write("")
                
                return copied_count
            
            # 复制文件并显示进度
            self.status_var.set("正在复制文件...")
            
            train_copied = copy_files(train_files, 'train')
            valid_copied = copy_files(valid_files, 'valid')
            test_copied = copy_files(test_files, 'test')
            
            # 更新data.yaml文件
            self.create_data_yaml(os.path.join(dataset_folder, "data.yaml"))
            
            # 显示完成信息
            result_msg = f"✅ 数据集划分完成！\n\n"
            result_msg += f"📊 划分结果：\n"
            result_msg += f"- 训练集: {train_count} 张 (包含标注: {train_copied} 张)\n"
            result_msg += f"- 验证集: {valid_count} 张 (包含标注: {valid_copied} 张)\n"
            result_msg += f"- 测试集: {test_count} 张 (包含标注: {test_copied} 张)\n\n"
            result_msg += f"📁 数据集保存位置: {dataset_folder}\n\n"
            result_msg += f"🎉 数据已准备好用于YOLO训练！"
            
            messagebox.showinfo("划分完成", result_msg)
            self.status_var.set("数据集划分完成")
            
        except Exception as e:
            messagebox.showerror("错误", f"划分数据集失败: {str(e)}")
            self.status_var.set("数据集划分失败")
    
    def show_annotation_guide(self):
        """显示标注指南"""
        guide_text = """
📖 YOLO数据标注指南

🎯 标注目标：
为每个图像中的目标物体绘制边界框或多边形，用于训练YOLO模型。

📋 标注步骤：

1️⃣ 添加类别：
   • 在"类别管理"区域输入类别名称
   • 点击"添加"按钮
   • 例如：电瓶车、人、电梯等

2️⃣ 选择标注模式：
   • 矩形框模式：适合规则形状的物体
   • 多边形模式：适合不规则形状的物体

3️⃣ 开始标注：
   • 选择要标注的类别
   • 在图像上绘制边界框或多边形
   • 点击"确认矩形"或"完成多边形"

4️⃣ 保存标注：
   • 按Ctrl+S或点击"保存标注"
   • 标注会自动保存为.txt文件

⚠️ 重要注意事项：

✅ 坐标范围：所有坐标会自动归一化到0-1之间
✅ 类别ID：从0开始，按添加顺序递增
✅ 文件格式：每行一个目标，格式为：类别ID 坐标1 坐标2 ...

❌ 常见错误：
• 坐标超出0-1范围（软件会自动处理）
• 类别ID超出范围（检查类别数量）
• 标注文件格式错误（使用软件标注避免）

🔧 快捷键：
• Delete：删除选中标注
• Escape：取消当前绘制
• Ctrl+S：保存标注
• ←→：上一张/下一张图像

💡 提示：
• 标注要准确覆盖目标物体
• 避免标注背景或无关区域
• 保持标注的一致性
• 定期保存标注文件
        """
        
        # 创建指南窗口
        guide_window = tk.Toplevel(self.root)
        guide_window.title("标注指南")
        guide_window.geometry("600x700")
        guide_window.transient(self.root)
        guide_window.grab_set()
        
        # 创建滚动文本框
        text_frame = ttk.Frame(guide_window)
        text_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        text_widget = tk.Text(text_frame, wrap=tk.WORD, font=("Arial", 10))
        scrollbar = ttk.Scrollbar(text_frame, orient=tk.VERTICAL, command=text_widget.yview)
        text_widget.configure(yscrollcommand=scrollbar.set)
        
        text_widget.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        text_widget.insert(tk.END, guide_text)
        text_widget.config(state=tk.DISABLED)
        
        # 关闭按钮
        ttk.Button(guide_window, text="关闭", command=guide_window.destroy).pack(pady=10)
    
    def show_welcome_message(self):
        """显示欢迎消息和快速指南"""
        welcome_text = f"""
🎉 欢迎使用YOLO数据标注工具！

📁 已加载文件夹：{os.path.basename(self.image_folder)}
📊 图像数量：{len(self.image_files)} 张
🏷️ 类别数量：{len(self.classes)} 个

🚀 快速开始：

1️⃣ 添加类别（如果还没有）：
   • 在左侧"类别管理"区域输入类别名称
   • 点击"添加"按钮

2️⃣ 开始标注：
   • 选择要标注的类别
   • 在图像上绘制边界框或多边形
   • 按Ctrl+S保存标注

3️⃣ 验证数据：
   • 点击"验证数据"检查标注质量
   • 查看"标注指南"了解详细说明

💡 提示：点击"标注指南"查看完整使用说明
        """
        
        messagebox.showinfo("欢迎使用", welcome_text)
    
    def export_dataset(self):
        """导出数据集"""
        if not hasattr(self, 'image_folder') or not self.image_folder:
            messagebox.showwarning("警告", "请先选择图像文件夹")
            return
        
        # 选择导出格式
        format_dialog = tk.Toplevel(self.root)
        format_dialog.title("选择导出格式")
        format_dialog.geometry("300x200")
        format_dialog.transient(self.root)
        format_dialog.grab_set()
        
        ttk.Label(format_dialog, text="选择导出格式:").pack(pady=10)
        
        format_var = tk.StringVar(value="yolo")
        ttk.Radiobutton(format_dialog, text="YOLO格式", variable=format_var, value="yolo").pack(pady=5)
        ttk.Radiobutton(format_dialog, text="COCO格式", variable=format_var, value="coco").pack(pady=5)
        ttk.Radiobutton(format_dialog, text="Pascal VOC格式", variable=format_var, value="voc").pack(pady=5)
        
        def export_selected():
            format_type = format_var.get()
            format_dialog.destroy()
            
            if format_type == "yolo":
                messagebox.showinfo("信息", "YOLO格式已保存")
            elif format_type == "coco":
                messagebox.showinfo("信息", "COCO格式导出功能待实现")
            elif format_type == "voc":
                messagebox.showinfo("信息", "Pascal VOC格式导出功能待实现")
        
        ttk.Button(format_dialog, text="导出", command=export_selected).pack(pady=20)
    
    # 快捷键功能
    def quick_select_class(self, class_index):
        """快速选择类别"""
        if 0 <= class_index < len(self.classes):
            self.class_var.set(self.classes[class_index])
            self.current_class = class_index
            self.status_var.set(f"已选择类别: {self.classes[class_index]}")
    
    def cancel_current_annotation(self):
        """取消当前标注"""
        if self.drawing:
            self.drawing = False
            if self.rect_id:
                self.canvas.delete(self.rect_id)
                self.rect_id = None
            self.status_var.set("已取消当前标注")
    
    def delete_selected_annotation(self):
        """删除选中的标注"""
        selection = self.annotation_tree.selection()
        if not selection:
            messagebox.showwarning("警告", "请先选择要删除的标注")
            return
        
        item = selection[0]
        index = int(self.annotation_tree.item(item, 'values')[0])
        
        if 0 <= index < len(self.annotations):
            del self.annotations[index]
            self.update_annotation_list()
            self.draw_annotations()
            self.status_var.set("已删除选中标注")
    
    def get_class_name(self, class_id):
        """根据标签ID获取标签名称"""
        if class_id < len(self.classes) and self.classes[class_id]:
            return self.classes[class_id]
        else:
            return f"class_{class_id}"
    
    def get_class_color(self, class_id):
        """根据标签ID获取标签颜色"""
        if class_id < len(self.class_colors):
            return self.class_colors[class_id]
        else:
            # 如果标签ID超出预定义颜色范围，使用循环颜色
            return self.class_colors[class_id % len(self.class_colors)]
    
    def rename_images_to_numbers(self, folder_path, image_files):
        """将图像文件重命名为数字编号"""
        try:
            import shutil
            
            # 创建临时文件夹用于重命名
            temp_folder = os.path.join(folder_path, "temp_rename")
            os.makedirs(temp_folder, exist_ok=True)
            
            # 先复制所有文件到临时文件夹，使用数字编号
            renamed_files = []
            for i, image_file in enumerate(image_files, 1):
                # 获取文件扩展名
                _, ext = os.path.splitext(image_file)
                
                # 生成新的文件名：00001.jpg, 00002.jpg, ...
                new_name = f"{i:05d}{ext}"
                new_path = os.path.join(temp_folder, new_name)
                
                # 复制文件到临时文件夹
                old_path = os.path.join(folder_path, image_file)
                shutil.copy2(old_path, new_path)
                renamed_files.append((image_file, new_name))
            
            # 删除原始文件
            for image_file in image_files:
                old_path = os.path.join(folder_path, image_file)
                if os.path.exists(old_path):
                    os.remove(old_path)
            
            # 将重命名后的文件移回原文件夹
            for old_name, new_name in renamed_files:
                temp_path = os.path.join(temp_folder, new_name)
                final_path = os.path.join(folder_path, new_name)
                shutil.move(temp_path, final_path)
            
            # 删除临时文件夹
            os.rmdir(temp_folder)
            
            # 同时重命名对应的标注文件
            for old_name, new_name in renamed_files:
                old_annotation = os.path.join(folder_path, os.path.splitext(old_name)[0] + ".txt")
                new_annotation = os.path.join(folder_path, os.path.splitext(new_name)[0] + ".txt")
                
                if os.path.exists(old_annotation):
                    shutil.move(old_annotation, new_annotation)
            
            print(f"已重命名 {len(image_files)} 个图像文件为数字编号")
            
        except Exception as e:
            print(f"重命名图像文件时出错: {e}")
            messagebox.showerror("错误", f"重命名图像文件失败: {str(e)}")
    
    def update_stats(self):
        """更新统计信息"""
        total_annotations = len(self.annotations)
        self.stats_var.set(f"总标注: {total_annotations}")
            
    def open_image(self):
        """打开图像文件"""
        file_path = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp *.tiff"), ("所有文件", "*.*")]
        )
        if file_path:
            self.load_image(file_path)
            
    def open_folder(self):
        """打开图像文件夹"""
        folder_path = filedialog.askdirectory(title="选择图像文件夹")
        if folder_path:
            self.create_dataset_structure(folder_path)
            self.load_folder(folder_path)
            
    def load_image(self, image_path):
        """加载图像到画布"""
        try:
            # 先释放之前的图像资源
            if self.current_image:
                self.current_image.close()
                self.current_image = None
            if hasattr(self, 'display_image') and self.display_image:
                self.display_image.close()
                self.display_image = None
            if self.photo:
                self.photo = None
            
            self.current_image_path = image_path
            self.current_image = Image.open(image_path)
            
            # 计算缩放比例以适应画布
            canvas_width = self.canvas.winfo_width()
            canvas_height = self.canvas.winfo_height()
            
            if canvas_width <= 1 or canvas_height <= 1:
                # 如果画布还没有初始化，使用默认大小
                canvas_width, canvas_height = 800, 600
                
            img_width, img_height = self.current_image.size
            
            # 计算缩放比例，保持宽高比
            scale_x = canvas_width / img_width
            scale_y = canvas_height / img_height
            self.scale = min(scale_x, scale_y, 1.0)  # 不放大图像
            
            # 计算显示尺寸
            self.display_width = int(img_width * self.scale)
            self.display_height = int(img_height * self.scale)
            
            # 缩放图像
            self.display_image = self.current_image.resize((self.display_width, self.display_height), Image.LANCZOS)
            self.photo = ImageTk.PhotoImage(self.display_image)
            
            # 清除画布并显示图像
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)
            
            # 如果还没有加载标签信息，尝试从当前图像所在文件夹加载
            if not self.classes and os.path.dirname(image_path):
                image_folder = os.path.dirname(image_path)
                if image_folder != self.image_folder:
                    self.image_folder = image_folder
                    self.load_classes_file()
            
            # 加载对应的标注文件
            self.load_annotations_for_image()
            
            self.status_var.set(f"已加载: {os.path.basename(image_path)}")
            
        except Exception as e:
            messagebox.showerror("错误", f"无法加载图像: {str(e)}")
            
    def load_annotations_for_image(self):
        """加载当前图像的标注"""
        if not self.current_image_path:
            return
            
        # 查找对应的标注文件
        base_name = os.path.splitext(self.current_image_path)[0]
        annotation_file = base_name + ".txt"
        
        self.annotations = []
        if os.path.exists(annotation_file):
            try:
                with open(annotation_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            class_id = int(parts[0])
                            
                            # 判断是矩形还是多边形
                            if len(parts) == 5:  # 矩形格式：class_id x_center y_center width height
                                x_center = float(parts[1])
                                y_center = float(parts[2])
                                width = float(parts[3])
                                height = float(parts[4])
                                
                                # 转换为像素坐标
                                img_width, img_height = self.current_image.size
                                x1 = (x_center - width/2) * img_width
                                y1 = (y_center - height/2) * img_height
                                x2 = (x_center + width/2) * img_width
                                y2 = (y_center + height/2) * img_height
                                
                                self.annotations.append({
                                    'class_id': class_id,
                                    'class_name': self.get_class_name(class_id),
                                    'bbox': (x1, y1, x2, y2),
                                    'type': 'rectangle'
                                })
                                
                            else:  # 多边形格式：class_id x1 y1 x2 y2 ... xn yn
                                # 获取图像尺寸
                                img_width, img_height = self.current_image.size
                                points = []
                                for i in range(1, len(parts), 2):
                                    if i + 1 < len(parts):
                                        x = float(parts[i]) * img_width
                                        y = float(parts[i + 1]) * img_height
                                        points.append((x, y))
                                
                                if len(points) >= 3:
                                    self.annotations.append({
                                        'class_id': class_id,
                                        'class_name': self.get_class_name(class_id),
                                        'points': points,
                                        'type': 'polygon'
                                    })
            except Exception as e:
                messagebox.showerror("错误", f"无法加载标注文件: {str(e)}")
                
        self.update_annotation_list()
        self.draw_annotations()
        self.update_stats()
        self.update_stats()
        
    def on_mode_change(self):
        """标注模式切换"""
        self.annotation_mode = self.mode_var.get()
        if self.annotation_mode == "polygon":
            self.canvas.configure(cursor="crosshair")
            self.status_var.set("多边形模式：左键添加点，右键完成")
        else:
            self.canvas.configure(cursor="crosshair")
            self.status_var.set("矩形框模式：拖拽绘制矩形")
            
    def on_mouse_press(self, event):
        """鼠标按下事件"""
        if not self.current_image:
            return
            
        if self.annotation_mode == "rectangle":
            self.drawing = True
            self.start_x = event.x
            self.start_y = event.y
            # 删除当前正在绘制的矩形
            self.canvas.delete("current_bbox")
        elif self.annotation_mode == "polygon":
            # 添加多边形点
            self.polygon_points.append((event.x, event.y))
            self.draw_polygon_preview()
        
    def on_mouse_drag(self, event):
        """鼠标拖拽事件"""
        if not self.drawing or not self.current_image or self.annotation_mode != "rectangle":
            return
            
        # 删除之前的临时矩形
        self.canvas.delete("current_bbox")
        
        # 获取当前标签的颜色
        color = self.get_class_color(self.current_class) if self.current_class >= 0 else "#FF0000"
        
        # 绘制新的临时矩形
        self.canvas.create_rectangle(
            self.start_x, self.start_y, event.x, event.y,
            outline=color, width=2, tags="current_bbox"
        )
        
    def on_mouse_release(self, event):
        """鼠标释放事件"""
        if not self.drawing or not self.current_image or self.annotation_mode != "rectangle":
            return
            
        self.drawing = False
        
        # 删除临时矩形
        self.canvas.delete("current_bbox")
        
        # 计算矩形坐标
        x1, y1 = min(self.start_x, event.x), min(self.start_y, event.y)
        x2, y2 = max(self.start_x, event.x), max(self.start_y, event.y)
        
        # 检查矩形大小
        if abs(x2 - x1) < 5 or abs(y2 - y1) < 5:
            return
            
        # 保存当前矩形到current_bbox，等待确认
        self.current_bbox = (x1, y1, x2, y2)
        
        # 获取当前标签的颜色
        color = self.get_class_color(self.current_class) if self.current_class >= 0 else "#FF0000"
        
        # 绘制确认矩形
        self.rect_id = self.canvas.create_rectangle(x1, y1, x2, y2, outline=color, width=2, tags="current_bbox")
        
        self.status_var.set("矩形已绘制，请点击'确认矩形'按钮完成标注")
        
    def on_right_click(self, event):
        """右键完成多边形"""
        if self.annotation_mode == "polygon" and len(self.polygon_points) >= 3:
            self.finish_polygon()
            
    def draw_polygon_preview(self):
        """绘制多边形预览"""
        if len(self.polygon_points) < 2:
            return
            
        # 删除之前的多边形预览
        self.canvas.delete("polygon_preview")
        
        # 获取当前标签的颜色
        color = self.get_class_color(self.current_class) if self.current_class >= 0 else "#FF0000"
        
        # 绘制多边形预览
        points = []
        for x, y in self.polygon_points:
            points.extend([x, y])
            
        if len(points) >= 6:  # 至少3个点
            self.canvas.create_polygon(points, outline=color, width=2, 
                                     fill="", tags="polygon_preview")
            
        # 绘制点
        for x, y in self.polygon_points:
            self.canvas.create_oval(x-3, y-3, x+3, y+3, fill=color, 
                                  outline=color, tags="polygon_preview")
            
    def confirm_rectangle(self):
        """确认矩形标注"""
        if not self.current_bbox:
            messagebox.showwarning("警告", "请先绘制矩形")
            return
        
        if self.current_class < 0 or self.current_class >= len(self.classes):
            messagebox.showwarning("警告", "请先选择类别")
            return
        
        # 转换为图像坐标
        img_bbox = [
            self.current_bbox[0] / self.scale,
            self.current_bbox[1] / self.scale,
            self.current_bbox[2] / self.scale,
            self.current_bbox[3] / self.scale
        ]
        
        # 创建矩形标注
        annotation = {
            'class_id': self.current_class,
            'class_name': self.classes[self.current_class],
            'bbox': img_bbox,
            'type': 'rectangle'
        }
        self.annotations.append(annotation)
        
        # 清理当前矩形
        self.current_bbox = None
        if self.rect_id:
            self.canvas.delete(self.rect_id)
            self.rect_id = None
        
        # 更新界面
        self.update_annotation_list()
        self.draw_annotations()
        self.update_stats()
        self.status_var.set("矩形标注完成")
    
    def finish_polygon(self):
        """完成多边形标注"""
        if len(self.polygon_points) < 3:
            messagebox.showwarning("警告", "多边形至少需要3个点")
            return
            
        # 转换为图像坐标
        img_points = []
        for x, y in self.polygon_points:
            img_x = x / self.scale
            img_y = y / self.scale
            img_points.append((img_x, img_y))
            
        # 添加多边形标注
        annotation = {
            'class_id': self.current_class,
            'class_name': self.classes[self.current_class],
            'points': img_points,
            'type': 'polygon'
        }
        self.annotations.append(annotation)
        
        # 清理多边形点
        self.polygon_points = []
        self.canvas.delete("polygon_preview")
        
        # 更新界面
        self.update_annotation_list()
        self.draw_annotations()
        self.update_stats()
        self.status_var.set("多边形标注完成")
        
    def draw_annotations(self):
        """绘制所有标注"""
        if not self.current_image:
            return
            
        # 清除之前的标注
        self.canvas.delete("annotation")
        
        for i, annotation in enumerate(self.annotations):
            class_name = annotation['class_name']
            class_id = annotation['class_id']
            annotation_type = annotation.get('type', 'rectangle')
            
            # 获取标签对应的颜色
            color = self.get_class_color(class_id)
            
            if annotation_type == 'rectangle':
                bbox = annotation['bbox']
                # 转换为显示坐标
                x1 = bbox[0] * self.scale
                y1 = bbox[1] * self.scale
                x2 = bbox[2] * self.scale
                y2 = bbox[3] * self.scale
                
                # 绘制矩形框
                rect_id = self.canvas.create_rectangle(
                    x1, y1, x2, y2,
                    outline=color, width=2, tags=("annotation", f"annotation_{i}")
                )
                
                # 绘制类别标签
                label_id = self.canvas.create_text(
                    x1, y1-10, anchor=tk.NW,
                    text=class_name, fill=color, font=("Arial", 10, "bold"),
                    tags=("annotation", f"annotation_{i}")
                )
                
            elif annotation_type == 'polygon':
                points = annotation['points']
                # 转换为显示坐标
                display_points = []
                for x, y in points:
                    display_x = x * self.scale
                    display_y = y * self.scale
                    display_points.extend([display_x, display_y])
                
                if len(display_points) >= 6:  # 至少3个点
                    # 绘制多边形
                    poly_id = self.canvas.create_polygon(
                        display_points, outline=color, width=2, fill="",
                        tags=("annotation", f"annotation_{i}")
                    )
                    
                    # 绘制类别标签（在多边形中心）
                    if len(points) > 0:
                        center_x = sum(p[0] for p in points) / len(points) * self.scale
                        center_y = sum(p[1] for p in points) / len(points) * self.scale
                        label_id = self.canvas.create_text(
                            center_x, center_y, anchor=tk.CENTER,
                            text=class_name, fill=color, font=("Arial", 10, "bold"),
                            tags=("annotation", f"annotation_{i}")
                        )
            
    def update_annotation_list(self):
        """更新标注列表"""
        # 清除现有项目
        for item in self.annotation_tree.get_children():
            self.annotation_tree.delete(item)
            
        # 添加标注项目
        for i, annotation in enumerate(self.annotations):
            annotation_type = annotation.get('type', 'rectangle')
            
            if annotation_type == 'rectangle':
                bbox = annotation['bbox']
                coord_str = f"({int(bbox[0])},{int(bbox[1])},{int(bbox[2])},{int(bbox[3])})"
            elif annotation_type == 'polygon':
                points = annotation['points']
                coord_str = f"多边形({len(points)}个点)"
            else:
                coord_str = "未知类型"
                
            self.annotation_tree.insert("", "end", values=(
                i,
                annotation['class_name'],
                coord_str
            ))
            
    def on_class_change(self, event):
        """类别选择改变"""
        selected_class = self.class_var.get()
        if selected_class in self.classes:
            self.current_class = self.classes.index(selected_class)
        elif selected_class == "请先添加类别":
            self.current_class = -1
            
    def add_class(self):
        """添加新类别"""
        new_class = self.new_class_var.get().strip()
        if new_class and new_class not in self.classes:
            self.classes.append(new_class)
            self.class_combo['values'] = self.classes
            self.class_var.set(new_class)
            self.current_class = len(self.classes) - 1
            self.new_class_var.set("")
            self.status_var.set(f"已添加类别: {new_class}")
            
            # 保存标签信息到文件
            self.save_classes_file()
            # 更新data.yaml文件
            self.update_data_yaml()
            # 更新统计信息
            self.update_stats()
        elif new_class in self.classes:
            messagebox.showwarning("警告", "类别已存在")
        else:
            messagebox.showwarning("警告", "请输入有效的类别名称")
            
    def on_annotation_select(self, event):
        """标注选择事件"""
        selection = self.annotation_tree.selection()
        if selection:
            item = self.annotation_tree.item(selection[0])
            annotation_id = int(item['values'][0])
            self.highlight_annotation(annotation_id)
            
    def highlight_annotation(self, annotation_id):
        """高亮显示选中的标注"""
        # 清除之前的高亮
        self.canvas.delete("highlight")
        
        if 0 <= annotation_id < len(self.annotations):
            annotation = self.annotations[annotation_id]
            annotation_type = annotation.get('type', 'rectangle')
            
            if annotation_type == 'rectangle':
                bbox = annotation['bbox']
                x1 = bbox[0] * self.scale
                y1 = bbox[1] * self.scale
                x2 = bbox[2] * self.scale
                y2 = bbox[3] * self.scale
                self.canvas.create_rectangle(
                    x1, y1, x2, y2,
                    outline="yellow", width=3, tags="highlight"
                )
            elif annotation_type == 'polygon':
                points = annotation['points']
                display_points = []
                for x, y in points:
                    display_x = x * self.scale
                    display_y = y * self.scale
                    display_points.extend([display_x, display_y])
                if len(display_points) >= 6:
                    self.canvas.create_polygon(
                        display_points, outline="yellow", width=3, fill="",
                        tags="highlight"
                    )
            
    def delete_selected(self):
        """删除选中的标注"""
        selection = self.annotation_tree.selection()
        if selection:
            item = self.annotation_tree.item(selection[0])
            annotation_id = int(item['values'][0])
            
            if 0 <= annotation_id < len(self.annotations):
                del self.annotations[annotation_id]
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set("已删除选中标注")
        else:
            messagebox.showwarning("警告", "请先选择要删除的标注")
            
    def clear_annotations(self):
        """清空所有标注"""
        if self.annotations:
            if messagebox.askyesno("确认", "确定要清空所有标注吗？"):
                self.annotations = []
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set("已清空所有标注")
        else:
            messagebox.showinfo("提示", "没有标注需要清空")
            
    def save_annotations(self):
        """保存标注为YOLO格式"""
        if not self.current_image_path or not self.annotations:
            messagebox.showwarning("警告", "没有标注需要保存")
            return
            
        try:
            # 生成标注文件路径
            base_name = os.path.splitext(self.current_image_path)[0]
            annotation_file = base_name + ".txt"
            
            # 获取图像尺寸
            img_width, img_height = self.current_image.size
            
            # 写入YOLO格式标注
            with open(annotation_file, 'w', encoding='utf-8') as f:
                for annotation in self.annotations:
                    class_id = annotation['class_id']
                    annotation_type = annotation.get('type', 'rectangle')
                    
                    if annotation_type == 'rectangle':
                        bbox = annotation['bbox']
                        # 转换为YOLO格式 (归一化坐标)
                        x_center = (bbox[0] + bbox[2]) / 2 / img_width
                        y_center = (bbox[1] + bbox[3]) / 2 / img_height
                        width = (bbox[2] - bbox[0]) / img_width
                        height = (bbox[3] - bbox[1]) / img_height
                        
                        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
                        
                    elif annotation_type == 'polygon':
                        points = annotation['points']
                        # 多边形格式：class_id x1 y1 x2 y2 ... xn yn
                        line = f"{class_id}"
                        for x, y in points:
                            norm_x = x / img_width
                            norm_y = y / img_height
                            line += f" {norm_x:.6f} {norm_y:.6f}"
                        f.write(line + "\n")
                    
            self.status_var.set(f"已保存标注: {os.path.basename(annotation_file)}")
            messagebox.showinfo("成功", f"标注已保存到: {annotation_file}")
            
        except Exception as e:
            messagebox.showerror("错误", f"保存标注失败: {str(e)}")
            
    def load_annotations(self):
        """加载标注文件"""
        if not self.current_image_path:
            messagebox.showwarning("警告", "请先打开图像")
            return
            
        file_path = filedialog.askopenfilename(
            title="选择标注文件",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        
        if file_path:
            try:
                self.annotations = []
                with open(file_path, 'r', encoding='utf-8') as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            class_id = int(parts[0])
                            
                            # 判断是矩形还是多边形
                            if len(parts) == 5:  # 矩形格式：class_id x_center y_center width height
                                x_center = float(parts[1])
                                y_center = float(parts[2])
                                width = float(parts[3])
                                height = float(parts[4])
                                
                                # 转换为像素坐标
                                img_width, img_height = self.current_image.size
                                x1 = (x_center - width/2) * img_width
                                y1 = (y_center - height/2) * img_height
                                x2 = (x_center + width/2) * img_width
                                y2 = (y_center + height/2) * img_height
                                
                                self.annotations.append({
                                    'class_id': class_id,
                                    'class_name': self.get_class_name(class_id),
                                    'bbox': (x1, y1, x2, y2),
                                    'type': 'rectangle'
                                })
                                
                            else:  # 多边形格式：class_id x1 y1 x2 y2 ... xn yn
                                # 获取图像尺寸
                                img_width, img_height = self.current_image.size
                                points = []
                                for i in range(1, len(parts), 2):
                                    if i + 1 < len(parts):
                                        x = float(parts[i]) * img_width
                                        y = float(parts[i + 1]) * img_height
                                        points.append((x, y))
                                
                                if len(points) >= 3:
                                    self.annotations.append({
                                        'class_id': class_id,
                                        'class_name': self.get_class_name(class_id),
                                        'points': points,
                                        'type': 'polygon'
                                    })
                                
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set(f"已加载标注: {os.path.basename(file_path)}")
                
            except Exception as e:
                messagebox.showerror("错误", f"加载标注失败: {str(e)}")
                
    def on_key_press(self, event):
        """键盘事件处理"""
        if event.keysym == "Delete":
            self.delete_selected()
        elif event.keysym == "Escape":
            self.canvas.delete("current_bbox")
            self.canvas.delete("polygon_preview")
            self.drawing = False
            if self.annotation_mode == "polygon":
                self.polygon_points = []
                self.status_var.set("已取消多边形绘制")
        elif event.keysym == "s" and event.state & 0x4:  # Ctrl+S
            self.save_annotations()
        elif event.keysym == "o" and event.state & 0x4:  # Ctrl+O
            self.open_image()
        elif event.keysym == "Return" and self.annotation_mode == "polygon":  # Enter完成多边形
            self.finish_polygon()
        elif event.keysym == "Left" and self.image_files:  # 左箭头键上一张
            self.prev_image()
        elif event.keysym == "Right" and self.image_files:  # 右箭头键下一张
            self.next_image()
    
    def on_closing(self):
        """窗口关闭事件处理 - 释放所有资源并终止程序"""
        import sys
        
        try:
            # 保存当前标注（如果有未保存的，静默保存，不弹出对话框）
            if self.current_image_path and self.annotations:
                try:
                    self.save_current_annotations()
                except Exception as e:
                    print(f"保存标注时出错（已忽略）: {e}")
            
        except Exception as e:
            print(f"保存数据时出错（已忽略）: {e}")
        
        try:
            # 释放所有资源
            self.cleanup_resources()
        except Exception as e:
            print(f"清理资源时出错（已忽略）: {e}")
        
        try:
            # 销毁窗口并退出程序
            self.root.quit()
            self.root.destroy()
        except Exception as e:
            print(f"销毁窗口时出错（已忽略）: {e}")
        
        # 强制退出程序（确保所有资源释放和进程终止）
        sys.exit(0)
    
    def cleanup_resources(self):
        """释放所有资源"""
        try:
            # 清理图像对象（PIL Image对象）
            try:
                if self.current_image:
                    self.current_image.close()
                    self.current_image = None
            except Exception as e:
                print(f"清理current_image时出错: {e}")
            
            try:
                if hasattr(self, 'display_image') and self.display_image:
                    self.display_image.close()
                    self.display_image = None
            except Exception as e:
                print(f"清理display_image时出错: {e}")
            
            # 清理PhotoImage对象（tkinter会管理，但显式设置为None有助于垃圾回收）
            self.photo = None
            
            # 清理画布上的所有对象（不删除canvas本身，因为它是窗口的一部分）
            try:
                if hasattr(self, 'canvas') and self.canvas:
                    self.canvas.delete("all")
            except Exception as e:
                print(f"清理canvas时出错: {e}")
            
            # 清空所有列表和数据结构（释放内存）
            self.annotations.clear()
            self.classes.clear()
            self.image_files.clear()
            self.polygon_points.clear()
            self.undo_stack.clear()
            self.redo_stack.clear()
            
            # 清理临时变量和引用
            self.current_image_path = None
            self.current_bbox = None
            self.current_polygon = None
            self.image_folder = None
            self.current_class = 0
            self.drawing = False
            self.start_x = 0
            self.start_y = 0
            
            # 清理状态变量
            if hasattr(self, 'status_var'):
                try:
                    self.status_var.set("")
                except:
                    pass
            
            # 强制垃圾回收，释放内存
            import gc
            gc.collect()
            
            print("✅ 资源已释放")
            
        except Exception as e:
            print(f"清理资源时出错: {e}")


def create_yolo_annotator(folder_path=None):
    """
    创建YOLO标注工具界面
    
    Args:
        folder_path (str, optional): 图像文件夹路径，如果提供会自动加载文件夹中的图像
        
    Returns:
        YOLOAnnotator: 标注工具实例
    """
    root = tk.Tk()
    app = YOLOAnnotator(root, initial_folder=folder_path)
    return app


def run_yolo_annotator(folder_path=None):
    """
    运行YOLO标注工具
    
    Args:
        folder_path (str, optional): 图像文件夹路径，如果提供会自动加载文件夹中的图像
    """
    # 如果没有提供文件夹路径，尝试从环境变量获取
    if not folder_path:
        folder_path = os.environ.get('YOLO_ANNOTATION_FOLDER')
        if folder_path:
            print(f"从环境变量获取文件夹路径: {folder_path}")
    
    root = tk.Tk()
    app = YOLOAnnotator(root, initial_folder=folder_path)
    root.mainloop()


if __name__ == "__main__":
    # 示例用法
    import sys
    
    # 优先检查环境变量
    folder_path = os.environ.get('YOLO_ANNOTATION_FOLDER')
    
    if folder_path and os.path.exists(folder_path):
        print(f"从环境变量加载文件夹: {folder_path}")
        run_yolo_annotator(folder_path)
    elif len(sys.argv) > 1:
        # 如果提供了命令行参数，使用第一个参数作为文件夹路径
        folder_path = sys.argv[1]
        run_yolo_annotator(folder_path)
    else:
        # 没有参数，启动空界面
        run_yolo_annotator()
